"""
Energy scan script for HEMC validation.
Runs Geant4 simulations across multiple energy points.
"""

import os
import subprocess
import sys
import threading
from pathlib import Path
from typing import cast

from jinja2 import Environment, FileSystemLoader

from ..config.models import HEMCConfig
from ..config.router import router

_MODULENAME = Path(__file__).stem
## Setup Jinja2 template environment
# _TEMPLATE_DIR = Path(__file__).parent.parent / "templates"
# _ENV = Environment(loader=FileSystemLoader(_TEMPLATE_DIR))


def generate_energy_list(config: HEMCConfig) -> list:
    """
    Generate energy sampling list (in MeV).
    Returns:
        List of energy values in MeV
    """
    ## 优先使用 确定的能量点列表, 方便测试
    if config.energyscan.energy_list:
        return config.energyscan.energy_list
    else:
        stt = int(config.energyscan.e_min)
        end = int(config.energyscan.e_max)
        step = int(config.energyscan.e_step)
        iend = (end + step - 1) // step * step + step  # 包含 end
        return list(range(stt, iend, step))

    return [
        1.000e-03,
        5.000e-03,
        1.000e-02,
        5.000e-02,
        1.000e-01,
        5.000e-01,
        1.000e00,
        5.000e00,
        1.000e01,
        5.000e01,
        1.000e02,
        5.000e02,
        7.000e02,
        1.000e03,
    ]


def format_energy(energy: float) -> str:
    """
    Format energy value to string with 1 decimal place for consistent filenames.
    Args:
        energy: Energy value in MeV
    Returns:
        Formatted string representation
    """
    return f"{energy:.2f}"


def tee_thread(pipe, file_obj):
    """Pipe content to both file and stdout in real-time."""
    for line in iter(pipe.readline, b""):
        line_dec = line.decode(errors="ignore")
        file_obj.write(line_dec)
        file_obj.flush()
        sys.stdout.write(line_dec)
        sys.stdout.flush()
    pipe.close()


def run_with_tee(cmd_list, log_path):
    """
    Run command with real-time output to both terminal and log file.
    Args:
        cmd_list: Command and arguments as list
        log_path: Path to log file
    Returns:
        Return code from subprocess
    """
    with open(log_path, "w", encoding="utf-8") as log_file:
        proc = subprocess.Popen(cmd_list, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, bufsize=1)
        thread = threading.Thread(target=tee_thread, args=(proc.stdout, log_file), daemon=True)
        thread.start()
        ret = proc.wait()
        thread.join()
        return ret


def _fmt(value, spec):
    return format(value, spec)


def _lifmt(values, spec: str):
    """
    对 list[float] 里的每个元素做 format(value, spec)
    返回 list[str]，保持顺序不变。
    如果传进来的不是列表/元组，直接原样返回，避免模板报错。
    """
    if not isinstance(values, (list, tuple)):
        return values
    return " ".join(format(v, spec) for v in values)


@router.register()
def energyscan(config: HEMCConfig):
    """Main entry point for energy scan."""
    ## config = create_config_from_fingerprint(*args, **kwargs)
    ## 检查
    outdir = cast(Path, config.project_paths.g4outdir)
    macrosdir = cast(Path, config.project_paths.macrosdir)
    logdir = cast(Path, config.project_paths.logdir)
    g4exe_path = cast(Path, config.project_paths.g4exe_path)

    ## geant4 jinja2 macro
    g4template_p = cast(Path, config.geant4.g4macro_template)
    env = Environment(loader=FileSystemLoader(g4template_p.parent))
    # env.filters['fmt'] = _fmt
    # env.filters['lifmt'] = _lifmt
    g4template = env.get_template(g4template_p.name)

    # Change to output directory for execution
    original_cwd = Path.cwd()

    try:
        os.chdir(outdir)

        energies = generate_energy_list(config)
        print(f"即将扫描 {len(energies)} 个能量点: ")
        print([format_energy(e) for e in energies])

        for engy in energies:
            estr = format_energy(engy)
            mac_file = macrosdir / f"letd_{estr}MeV.mac"
            log_file = logdir / f"log_{estr}MeV.log"

            ## 修改能量点; Write macro file using Jinja2 template
            g4cfg = config.geant4.to_template_dict(engy)
            macro_content = g4template.render(**g4cfg)
            ## ic(macro_content)

            with open(mac_file, "w", encoding="utf-8") as f:
                f.write(macro_content)

            # Run Geant4 simulation
            cmd_list = [g4exe_path, mac_file]
            ret = run_with_tee(cmd_list, log_file)
            if ret != 0:
                print(f"!!! 能量 {estr} MeV 运行失败, returncode = {ret}")

    finally:
        # Always restore original working directory
        os.chdir(original_cwd)

    print(f"\n全部扫描完成! 结果保存在目录: {config.project_paths.g4outdir}")


if __name__ == "__main__":
    pass
